{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4185a71b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "from torchvision.transforms import autoaugment\n",
    "from pathlib import Path\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "599e6d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据增强\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.RandomResizedCrop(224), # 随机裁剪\n",
    "    transforms.RandomHorizontalFlip(), # 随机水平翻转\n",
    "    transforms.RandomVerticalFlip(), # 随机垂直翻转\n",
    "    autoaugment.TrivialAugmentWide(),  # 自动数据增强\n",
    "    transforms.ToTensor(), # 转换为张量\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    transforms.RandomErasing(p=0.5, scale=(0.02, 0.1)),  # 随机擦除\n",
    "])\n",
    "val_transform = transforms.Compose([\n",
    "    transforms.Resize((256,256)),\n",
    "    transforms.CenterCrop(224), # 中心裁剪\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "143eb72e",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets_root = Path(\"fruit81_split\")  # 替换为你的实际路径\n",
    "train_data_path = datasets_root / \"train\"\n",
    "val_data_path = datasets_root / \"val\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ecdaf5a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = datasets.ImageFolder(root=train_data_path, transform=train_transform)\n",
    "val_dataset = datasets.ImageFolder(root=val_data_path, transform=val_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "88d21538",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset ImageFolder\n",
      "    Number of datapoints: 6997\n",
      "    Root location: fruit81_split\\val\n",
      "    StandardTransform\n",
      "Transform: Compose(\n",
      "               Resize(size=(256, 256), interpolation=bilinear, max_size=None, antialias=True)\n",
      "               CenterCrop(size=(224, 224))\n",
      "               ToTensor()\n",
      "               Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])\n",
      "           )\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = 32\n",
    "\n",
    "# 训练集的数据加载器\n",
    "train_loader = DataLoader(train_dataset,\n",
    "                          batch_size=BATCH_SIZE,\n",
    "                          shuffle=True,\n",
    "                          num_workers=1\n",
    "                         )\n",
    "\n",
    "# 测试集的数据加载器\n",
    "test_loader = DataLoader(val_dataset,\n",
    "                         batch_size=BATCH_SIZE,\n",
    "                         shuffle=False,\n",
    "                         num_workers=1\n",
    "                        )\n",
    "print(test_loader.dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0ab315d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 3, 224, 224])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "img, la = next(iter(train_loader))\n",
    "img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "3e4f5f4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. 构建 TensorFormer（简化版的 Vision Transformer）\n",
    "# 图片嵌入编码和word embedding一样都是给输入的数据进行高维映射\n",
    "\n",
    "class PatchEmbedding(nn.Module):\n",
    "    \"\"\"将图像分割为小块并嵌入为向量的模块（Vision Transformer的核心组件）\"\"\"\n",
    "    def __init__(self, in_channels=3, patch_size=16, emb_size=256, img_size=224):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            in_channels: 输入图像的通道数（RGB为3）\n",
    "            patch_size: 每个图像块的大小（默认16x16像素）\n",
    "            emb_size: 嵌入向量的维度\n",
    "            img_size: 输入图像的尺寸（假设为正方形）\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.patch_size = patch_size\n",
    "        \n",
    "        # 使用卷积操作实现分块嵌入（等效于将图像切块后线性投影）\n",
    "        self.proj = nn.Conv2d(in_channels, emb_size, \n",
    "                             kernel_size=patch_size, stride=patch_size)\n",
    "        \n",
    "        # 可学习的分类标记（CLS token），用于最终分类\n",
    "        self.cls_token = nn.Parameter(torch.randn(1, 1, emb_size))\n",
    "        \n",
    "        # 计算图像被分割后的块数（例如128x128图像分为(128/16)^2=64块）\n",
    "        num_patches = (img_size // patch_size) ** 2\n",
    "        \n",
    "        # 可学习的位置编码（+1是因为要加入CLS token）\n",
    "        self.pos_embed = nn.Parameter(torch.randn(1, num_patches + 1, emb_size))\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"输入: [B, C, H, W]  输出: [B, num_patches+1, emb_size]\"\"\"\n",
    "        B = x.shape[0]  # 批大小\n",
    "        \n",
    "        # 1. 分块投影 [B, 3, 128, 128] -> [B, 256, 8, 8]（假设patch_size=16）\n",
    "        x = self.proj(x)  \n",
    "        \n",
    "        # 2. 展平块维度 [B, 256, 8, 8] -> [B, 256, 64] -> [B, 64, 256]\n",
    "        x = x.flatten(2).transpose(1, 2)  \n",
    "        \n",
    "        # 3. 添加CLS token [1,1,256] -> [B,1,256] 并拼接\n",
    "        cls_tokens = self.cls_token.expand(B, -1, -1)\n",
    "        x = torch.cat([cls_tokens, x], dim=1)  # [B, 65, 256]\n",
    "        \n",
    "        # 4. 添加位置编码\n",
    "        x = x + self.pos_embed\n",
    "        return x\n",
    "\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    \"\"\"Transformer编码器层（包含多头注意力和前馈网络）\"\"\"\n",
    "    def __init__(self, emb_size=256, heads=8, forward_expansion=4, dropout=0.1):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            emb_size: 嵌入维度\n",
    "            heads: 注意力头的数量\n",
    "            forward_expansion: MLP隐藏层扩展倍数\n",
    "            dropout: Dropout率\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        # 层归一化（Pre-Norm结构）\n",
    "        self.ln1 = nn.LayerNorm(emb_size)\n",
    "        \n",
    "        # 多头注意力机制（batch_first=True表示输入为[B, N, D]）\n",
    "        self.attn = nn.MultiheadAttention(\n",
    "            embed_dim=emb_size, \n",
    "            num_heads=heads, \n",
    "            dropout=dropout, \n",
    "            batch_first=True\n",
    "        )\n",
    "        \n",
    "        self.ln2 = nn.LayerNorm(emb_size)\n",
    "        \n",
    "        # 前馈网络（MLP）\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(emb_size, emb_size * forward_expansion),  # 扩展维度\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(emb_size * forward_expansion, emb_size),  # 降回原维度\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"输入输出形状: [B, num_patches+1, emb_size]\"\"\"\n",
    "        # 1. 残差连接 + 多头注意力\n",
    "        attn_output, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x))\n",
    "        x = attn_output + x  # 残差连接\n",
    "        \n",
    "        # 2. 残差连接 + MLP\n",
    "        x = self.mlp(self.ln2(x)) + x\n",
    "        return x\n",
    "\n",
    "\n",
    "class TensorFormer(nn.Module):\n",
    "    \"\"\"完整的Transformer分类模型\"\"\"\n",
    "    def __init__(self, num_classes=81):\n",
    "        super().__init__()\n",
    "        # 1. 图像分块嵌入\n",
    "        self.patch_embed = PatchEmbedding()\n",
    "        \n",
    "        # 2. 堆叠4个Transformer编码器\n",
    "        self.encoder = nn.Sequential(*[TransformerEncoder() for _ in range(4)])\n",
    "        \n",
    "        # 3. 输出层前的归一化\n",
    "        self.ln = nn.LayerNorm(256)\n",
    "        \n",
    "        # 4. 分类头\n",
    "        self.head = nn.Linear(256, num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"输入: [B, C, H, W]  输出: [B, num_classes]\"\"\"\n",
    "        # 1. 分块嵌入 [B,3,128,128] -> [B,65,256]\n",
    "        x = self.patch_embed(x)  \n",
    "        \n",
    "        # 2. Transformer编码 [B,65,256] -> [B,65,256]\n",
    "        x = self.encoder(x)  \n",
    "        \n",
    "        # 3. 仅取CLS token对应的特征 [B,256]\n",
    "        x = self.ln(x[:, 0])  \n",
    "        \n",
    "        # 4. 分类预测\n",
    "        return self.head(x)  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "fdc1d136",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "81"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "num_class = len(os.listdir(train_data_path))\n",
    "num_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "999f37b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# 4. 初始化模型与训练参数\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = TensorFormer(num_classes=num_class).to(device) # 初始化模型\n",
    "\n",
    "criterion = nn.CrossEntropyLoss() # 交叉熵损失函数\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4) # Adam优化器\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6701f79b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from sklearn.metrics import precision_score\n",
    "from sklearn.metrics import recall_score\n",
    "from sklearn.metrics import accuracy_score\n",
    "from sklearn.metrics import f1_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "847b70a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. 训练循环\n",
    "\n",
    "\n",
    "def train_one_batch(images, labels, epoch, batch_idx):\n",
    "    '''\n",
    "    运行一个 batch 的训练，返回当前 batch 的训练日志\n",
    "    '''\n",
    "    \n",
    "    # 数据移动到 GPU\n",
    "    images = images.to(device)\n",
    "    labels = labels.to(device)\n",
    "    \n",
    "    outputs = model(images) # 输入模型，执行前向预测\n",
    "    loss = criterion(outputs, labels) # 计算当前 batch 中，每个样本的平均交叉熵损失函数值\n",
    "    \n",
    "    # 优化更新权重\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    # 返回的是预测结果，维度为[batch_size, num_classes]\n",
    "    # batch_size表示每一张图片，num_classes表示每个图片的类别数，里面放的是每个类别的概率，\n",
    "    # max是求最大值，dim = 1 表示在第 1 个维度上求最大值，也就是求每个num_classes的最大值\n",
    "    _, preds = torch.max(outputs, 1) \n",
    "    # 将预测值拷贝到 CPU，转换为 NumPy 数组\n",
    "    preds = preds.cpu().numpy()\n",
    "    # 将损失值拷贝到 CPU，转换为 NumPy 数组\n",
    "    # loss.detach()将损失值从计算图中分离出来，返回一个不需要梯度的新张量\n",
    "    loss = loss.detach().cpu().numpy()\n",
    "    # 将预测结果和真实标签拷贝到 CPU，转换为 NumPy 数组\n",
    "    outputs = outputs.detach().cpu().numpy()\n",
    "    labels = labels.detach().cpu().numpy()\n",
    "    \n",
    "    log_train = {}\n",
    "    log_train['epoch'] = epoch\n",
    "    log_train['batch'] = batch_idx\n",
    "    # 计算分类评估指标\n",
    "    log_train['train_loss'] = loss # 损失值\n",
    "    log_train['train_accuracy'] = accuracy_score(labels, preds) # 计算准确率\n",
    "    log_train['train_precision'] = precision_score(labels, preds, average='macro') # 计算精确率，，average='macro'表示计算所有类别的精确率，并取平均值\n",
    "    log_train['train_recall'] = recall_score(labels, preds, average='macro') # 计算召回率\n",
    "    log_train['train_f1-score'] = f1_score(labels, preds, average='macro') # 计算 F1-score\n",
    "    \n",
    "    return log_train\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "14f27bbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_testset(epoch):\n",
    "    '''\n",
    "    在整个测试集上评估，返回分类评估指标日志\n",
    "    '''\n",
    "\n",
    "    loss_list = []\n",
    "    labels_list = []\n",
    "    preds_list = []\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader: # 生成一个 batch 的数据和标注\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images) # 输入模型，执行前向预测\n",
    "\n",
    "            # 获取整个测试集的标签类别和预测类别\n",
    "            _, preds = torch.max(outputs, 1) # 获得当前 batch 所有图像的预测类别\n",
    "            \n",
    "            preds = preds.cpu().numpy()\n",
    "            loss = criterion(outputs, labels) # 由 logit，计算当前 batch 中，每个样本的平均交叉熵损失函数值\n",
    "            loss = loss.detach().cpu().numpy()\n",
    "            outputs = outputs.detach().cpu().numpy()\n",
    "            labels = labels.detach().cpu().numpy()\n",
    "\n",
    "            loss_list.append(loss)\n",
    "            labels_list.extend(labels)\n",
    "            preds_list.extend(preds)\n",
    "        \n",
    "    log_test = {}\n",
    "    log_test['epoch'] = epoch\n",
    "    \n",
    "    # 计算分类评估指标\n",
    "    log_test['test_loss'] = np.mean(loss_list)\n",
    "    log_test['test_accuracy'] = accuracy_score(labels_list, preds_list)\n",
    "    log_test['test_precision'] = precision_score(labels_list, preds_list, average='macro')\n",
    "    log_test['test_recall'] = recall_score(labels_list, preds_list, average='macro')\n",
    "    log_test['test_f1-score'] = f1_score(labels_list, preds_list, average='macro')\n",
    "    \n",
    "    return log_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "7b24955c",
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 0\n",
    "batch_idx = 0 # 批次索引\n",
    "best_test_accuracy = 0\n",
    "EPOCHS = 50\n",
    "df_train_log = pd.DataFrame()\n",
    "df_val_log = pd.DataFrame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "6f733eda",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:55<00:00,  3.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.161.pth\n",
      "Epoch 2/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:50<00:00,  4.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.211.pth\n",
      "Epoch 3/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:50<00:00,  4.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.251.pth\n",
      "Epoch 4/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.270.pth\n",
      "Epoch 5/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.289.pth\n",
      "Epoch 6/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.318.pth\n",
      "Epoch 7/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.339.pth\n",
      "Epoch 8/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.357.pth\n",
      "Epoch 9/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:51<00:00,  4.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.359.pth\n",
      "Epoch 10/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.386.pth\n",
      "Epoch 11/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.389.pth\n",
      "Epoch 12/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.403.pth\n",
      "Epoch 13/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.417.pth\n",
      "Epoch 14/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.424.pth\n",
      "Epoch 15/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.440.pth\n",
      "Epoch 16/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.450.pth\n",
      "Epoch 17/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:47<00:00,  4.15it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.471.pth\n",
      "Epoch 19/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.477.pth\n",
      "Epoch 20/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.496.pth\n",
      "Epoch 21/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.499.pth\n",
      "Epoch 22/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.503.pth\n",
      "Epoch 23/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:50<00:00,  4.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.518.pth\n",
      "Epoch 24/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.08it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.540.pth\n",
      "Epoch 25/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 26/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 27/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.554.pth\n",
      "Epoch 28/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.557.pth\n",
      "Epoch 29/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.579.pth\n",
      "Epoch 30/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 31/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 32/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.586.pth\n",
      "Epoch 33/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 34/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.599.pth\n",
      "Epoch 35/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.11it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 36/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.618.pth\n",
      "Epoch 37/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 38/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:50<00:00,  4.04it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 39/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.632.pth\n",
      "Epoch 40/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.641.pth\n",
      "Epoch 41/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.653.pth\n",
      "Epoch 42/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 43/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:47<00:00,  4.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 44/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.659.pth\n",
      "Epoch 45/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.669.pth\n",
      "Epoch 46/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:49<00:00,  4.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 47/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:48<00:00,  4.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 48/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [02:00<00:00,  3.73it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.683.pth\n",
      "Epoch 49/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:54<00:00,  3.92it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "保存新的最佳模型 checkpoint_1/best-0.687.pth\n",
      "Epoch 50/50\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 448/448 [01:52<00:00,  4.00it/s]\n"
     ]
    }
   ],
   "source": [
    "# 开始训练\n",
    "for epoch in range(1, EPOCHS+1): \n",
    "    \n",
    "    print(f'Epoch {epoch}/{EPOCHS}')\n",
    "    \n",
    "    ## 训练阶段，一批一批的训练数据\n",
    "    model.train()\n",
    "    for images, labels in tqdm(train_loader): # 获得一个 batch 的数据和标注\n",
    "        batch_idx += 1 # 训练次数加1\n",
    "        log_train = train_one_batch(images, labels, epoch, batch_idx) # 训练一个 batch 的数据\n",
    "        df_train_log = df_train_log._append(log_train, ignore_index=True) # 将训练日志添加到训练日志数据框中\n",
    "        \n",
    "#     # lr_scheduler.step() # 更新学习率\n",
    "\n",
    "    ## 测试阶段\n",
    "    model.eval()\n",
    "    log_test = evaluate_testset(epoch)\n",
    "    df_val_log = df_val_log._append(log_test, ignore_index=True)\n",
    "    \n",
    "    # 保存最新的最佳模型文件\n",
    "    if log_test['test_accuracy'] > best_test_accuracy: \n",
    "        # 删除旧的最佳模型文件(如有)\n",
    "        old_best_checkpoint_path = 'checkpoint_1/best-{:.3f}.pth'.format(best_test_accuracy)\n",
    "        if os.path.exists(old_best_checkpoint_path):\n",
    "            os.remove(old_best_checkpoint_path)\n",
    "        # 保存新的最佳模型文件\n",
    "        best_test_accuracy = log_test['test_accuracy']\n",
    "        new_best_checkpoint_path = 'checkpoint_1/best-{:.3f}.pth'.format(log_test['test_accuracy'])\n",
    "        torch.save(model.state_dict(), new_best_checkpoint_path)\n",
    "        print('保存新的最佳模型', 'checkpoint_1/best-{:.3f}.pth'.format(best_test_accuracy))\n",
    "        # best_test_accuracy = log_test['test_accuracy']\n",
    "\n",
    "df_train_log.to_csv(r'transformer_model/训练日志-训练集.csv', index=False)\n",
    "df_val_log.to_csv(r'transformer_model/训练日志-测试集.csv', index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "647cb310",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 7. 保存模型\n",
    "torch.save(model.state_dict(), \"transformer_model.pth\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
